from openai import OpenAI
import streamlit as st
from langchain.chains import RetrievalQA
from ChatGLM_new import tongyi_embeddings,sspy_llm
from langchain.prompts import PromptTemplate
from langchain.vectorstores import Chroma
def load_chain():
    # 加载问答链
    # 定义 Embeddings
    #embeddings = HuggingFaceEmbeddings(model_name="/root/data/model/sentence-transformer")
    embeddings = tongyi_embeddings

    # 向量数据库持久化路径
    persist_directory = 'data_base/vector_db/chroma'

    # 加载数据库
    vectordb = Chroma(
        persist_directory=persist_directory,  # 允许我们将persist_directory目录保存到磁盘上
        embedding_function=embeddings
    )

    #llm = InternLM_LLM(model_path="/root/data/model/Shanghai_AI_Laboratory/internlm-chat-7b")
    llm = sspy_llm
    # 你可以修改这里的 prompt template 来试试不同的问答效果
    template = """请使用以下提供的上下文来回答用户的问题。如果无法从上下文中得到答案，请回答你不知道，并总是使用中文回答。
    提供的上下文：
    ···
    {context}
    ···
    用户的问题: {question}
    你给的回答:"""

    QA_CHAIN_PROMPT = PromptTemplate(input_variables=["context", "question"],
                                     template=template)

    # 运行 chain
    qa_chain = RetrievalQA.from_chain_type(llm,
                                           retriever=vectordb.as_retriever(),
                                           return_source_documents=True,
                                           chain_type_kwargs={"prompt": QA_CHAIN_PROMPT})

    return qa_chain

class ModelCenter:
    """
    存储问答 Chain 的对象
    """
    def __init__(self):
        self.chain = load_chain()

    def qa_chain_self_answer(self, question: str, chat_history: list = []):
        """
        调用不带历史记录的问答链进行回答
        """
        if question == None or len(question) < 1:
            return "", chat_history
        try:
            result = self.chain({"query": question})["result"]
            chat_history.append(
                (question, result))
            return result, chat_history
        except Exception as e:
            return e, chat_history

model_center = ModelCenter()


with st.sidebar:
    openai_api_key = st.text_input("OpenAI API Key", key="chatbot_api_key", type="password")

st.title("💬 Chatbot")
st.caption("🚀 A chatbot powered by OpenAI")
if "messages" not in st.session_state:
    st.session_state["messages"] = [{"role": "assistant", "content": "我该怎么帮助你?"}]

for msg in st.session_state.messages:
    st.chat_message(msg["role"]).write(msg["content"])

# 创建一个聊天记录列表
chat_history = []
if prompt := st.chat_input():
    client = OpenAI(api_key=openai_api_key)
    st.session_state.messages.append({"role": "user", "content": prompt})
    st.chat_message("user").write(prompt)
    msg, chat_history = model_center.qa_chain_self_answer(prompt, chat_history)
    st.session_state.messages.append({"role": "assistant", "content": msg})
    st.chat_message("assistant").write(msg)